# This code is part of cqlib.
#
# (C) Copyright China Telecom Quantum Group, QuantumCTek Co., Ltd.,
# Center for Excellence in Quantum Information and Quantum Physics 2024.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.

"""
TianYanPlatform test case.

If you are in a test environment, you can set the environment variables
`SCHEME` and `DOMAIN` to direct to the test URL.
"""
# pylint: disable=redefined-outer-name,missing-function-docstring


import os
import random
from datetime import datetime

import pytest

from cqlib import TianYanPlatform
from cqlib.circuits import Circuit
from cqlib.quantum_platform.base import QuantumLanguage

login_key = os.getenv('LOGIN_KEY', '')
sim_machine = os.getenv('SIM_MACHINE', 'tensor_network')
quantum_machine = os.getenv('QUANTUM_MACHINE', 'xuanyuanone')


@pytest.fixture
def platform():
    """返回一个初始化的 TianYanPlatform 实例"""
    if scheme := os.getenv('SCHEME'):
        TianYanPlatform.SCHEME = scheme
    if domain := os.getenv('DOMAIN'):
        TianYanPlatform.DOMAIN = domain
    assert login_key, 'Please set the login key first'
    assert sim_machine, 'Please set the simulator code first'
    assert quantum_machine, 'Please set the quantum computer code first'

    return TianYanPlatform(login_key=login_key, machine_name=sim_machine)


@pytest.fixture
def lab_id(platform):  # pylint: disable=redefined-outer-name
    """创建一个新的实验室并返回其 ID"""
    name = f'lab.{datetime.now().strftime("%Y%m%d%H%M%S.%f")}.{random.randint(0, 100000)}'
    return platform.create_lab(name)


@pytest.fixture
def qcis():  # pylint: disable=redefined-outer-name
    num_qubits = 5
    circuit = Circuit(num_qubits)
    for i in range(num_qubits):
        circuit.h(i)
    for i in range(num_qubits - 1):
        circuit.cz(i, i + 1)
    for i in range(num_qubits):
        circuit.x(i)
    circuit.measure_all()
    return circuit.qcis


# pylint: disable=too-many-public-methods
class TestTianYanPlatform:
    """
    Test TianYanPlatform
    """

    isq_circuit = """import std;
qbit q[2];

unit main() {
    H(q[0]);
    H(q[1]);
    M(q[0]);
    M(q[1]);
}"""

    def test_login(self, platform):
        """
        test platform login

        Args:
            platform:

        Returns:
        """
        print(f'access_token: {platform.access_token}')
        assert platform.access_token

    def test_query_quantum_computer_list(self, platform):
        computer_list = platform.query_quantum_computer_list()
        print(f'computer_list: {computer_list}')
        assert isinstance(computer_list, list)

    def test_create_lab(self, platform):
        lab_id = platform.create_lab(f'lab.{datetime.now().strftime("%Y%m%d%H%M%S")}')
        print(f'lab_id: {lab_id}')
        assert lab_id

    def test_save_experiment_qcis(self, platform, lab_id, qcis):
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=qcis)
        print(f'exp_id: {exp_id}')
        assert exp_id

    def test_save_experiment_isq(self, platform, lab_id):
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=self.isq_circuit,
                                          language=QuantumLanguage.ISQ)
        print(f'exp_id: {exp_id}')
        assert exp_id

    def test_re_save_experiment(self, platform, lab_id, qcis):
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=qcis)
        print(f'exp_id: {exp_id}')
        assert exp_id
        re_exp_id = platform.save_experiment(lab_id=lab_id, circuit=qcis)
        print(f'exp_id: {re_exp_id}')
        assert re_exp_id

    def test_run_experiment_qcis(self, platform, lab_id, qcis):
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=qcis,
                                          language=QuantumLanguage.QCIS)
        print(f'exp_id: {exp_id}')
        query_id = platform.run_experiment(exp_id=exp_id, num_shots=3000)
        print(f'query_id: {query_id}')
        assert query_id

    def test_run_experiment_isq(self, platform, lab_id):
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=self.isq_circuit,
                                          language=QuantumLanguage.ISQ)
        print(f'exp_id: {exp_id}')
        query_id = platform.run_experiment(exp_id=exp_id, num_shots=3000)
        print(f'query_id: {query_id}')
        assert query_id

    def test_run_experiment_qcis_without_verify(self, platform, lab_id, qcis):
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=qcis,
                                          language=QuantumLanguage.QCIS)
        print(f'exp_id: {exp_id}')
        query_id = platform.run_experiment(exp_id=exp_id, num_shots=3000, is_verify=False)
        print(f'query_id: {query_id}')
        assert query_id

    def test_run_experiment_isq_without_verify(self, platform, lab_id):
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=self.isq_circuit,
                                          language=QuantumLanguage.ISQ)
        print(f'exp_id: {exp_id}')
        query_id = platform.run_experiment(exp_id=exp_id, num_shots=3000, is_verify=False)
        print(f'query_id: {query_id}')
        assert query_id

    def test_run_experiment_larger_num_shots(self, platform, lab_id, qcis):
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=qcis)
        print(f'exp_id: {exp_id}')
        query_id = platform.run_experiment(exp_id=exp_id, num_shots=100000)
        print(f'query_id: {query_id}')
        assert query_id

    def test_submit_job(self, platform, lab_id, qcis):
        query_id = platform.submit_job(
            circuit=qcis,
            exp_name=f'exp.{datetime.now().strftime("%Y%m%d%H%M%S")}',
            lab_id=lab_id,
            num_shots=3000,
        )
        print(f'query_id: {query_id}')
        assert query_id

    def test_submit_job_isq_exp_circuit(self, platform, lab_id):
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=self.isq_circuit,
                                          language=QuantumLanguage.ISQ)
        print(f'exp_id: {exp_id}')
        query_id = platform.submit_job(
            circuit=self.isq_circuit,
            exp_id=exp_id,
            language=QuantumLanguage.ISQ,
            num_shots=3000,
        )
        print(f'query_id: {query_id}')
        assert query_id

    def test_submit_job_exp_circuit(self, platform, lab_id, qcis):
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=qcis,
                                          language=QuantumLanguage.QCIS)
        print(f'exp_id: {exp_id}')
        query_id = platform.submit_job(
            circuit=qcis,
            exp_id=exp_id,
            num_shots=3000,
        )
        print(f'query_id: {query_id}')
        assert query_id

    def test_submit_job_parameters(self, platform, lab_id):
        circuit = "RX Q0 {theta}\nM Q0"
        query_id = platform.submit_job(
            circuit=circuit,
            parameters=[['theta']],
            values=[[1.2]],
            exp_name=f'exp.{datetime.now().strftime("%Y%m%d%H%M%S")}',
            lab_id=lab_id,
            num_shots=3000,
        )
        print(f'query_id: {query_id}')
        assert query_id

    def test_submit_job_parameters_isq(self, platform, lab_id):
        query_id = platform.submit_job(
            circuit=[self.isq_circuit],
            exp_name=f'exp.{datetime.now().strftime("%Y%m%d%H%M%S")}',
            lab_id=lab_id,
            language=QuantumLanguage.ISQ,
            num_shots=3000,
        )
        print(f'query_id: {query_id}')
        assert query_id

    def test_submit_job_exp_id(self, platform, lab_id, qcis):
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=qcis,
                                          language=QuantumLanguage.QCIS)
        query_id = platform.submit_job(
            exp_id=exp_id,
            num_shots=3000,
        )
        print(f'query_id: {query_id}')
        assert query_id

    def test_submit_job_isq_2(self, platform):
        query_id = platform.submit_job(
            circuit=[self.isq_circuit],
            num_shots=3000,
            language=QuantumLanguage.ISQ,
        )
        print(f'query_id: {query_id}')
        assert query_id

    def test_submit_job_lab_exp(self, platform, lab_id, qcis):
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=qcis,
                                          language=QuantumLanguage.QCIS)
        print(f'exp_id: {exp_id}')
        query_id = platform.submit_job(
            exp_id=exp_id,
            lab_id=lab_id,
            num_shots=100,
            language=QuantumLanguage.QCIS,
        )
        print(f'query_id: {query_id}')
        assert query_id

    def test_query_experiment(self, platform, lab_id, qcis):
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=qcis,
                                          language=QuantumLanguage.QCIS)
        print(f'exp_id: {exp_id}')
        query_id = platform.run_experiment(exp_id=exp_id, num_shots=3000)
        print(f'query_id: {query_id}')
        exp = platform.query_experiment(query_id)
        print(f'exp: {exp}')
        assert exp

    def test_query_experiment_isq(self, platform, lab_id):
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=self.isq_circuit,
                                          language=QuantumLanguage.ISQ)
        print(f'exp_id: {exp_id}')
        query_id = platform.run_experiment(exp_id=exp_id, num_shots=3000)
        print(f'query_id: {query_id}')
        exp = platform.query_experiment(query_id)
        print(f'exp: {exp}')
        assert exp

    def test_download_config(self, platform):
        cfg = platform.download_config()
        print(f'cfg: {cfg}')
        assert cfg

    def test_qcis_check_regular(self, platform):
        data = platform.qcis_check_regular('H Q0\nH Q2\nCZ Q0 Q2\nM Q0\nM Q2')
        print(f'qcis_check_regular: {data}')
        assert data

    def test_qcis_check_regular_false(self, platform):
        data = platform.qcis_check_regular('HQ0\nH Q2\nCZ Q0 Q2\nM Q0\nM Q2')
        print(f'data:{data}')
        assert data is False

    def test_qcis_check_regular_quantum_machine_false_circuit(self, platform):
        platform.set_machine(quantum_machine)
        data = platform.qcis_check_regular('HQ0\nH Q2\nCZ Q0 Q2\nM Q0\nM Q2')
        platform.set_machine(sim_machine)
        print(f'data: {data}')
        assert data is False

    def test_qcis_check_regular_quantum_machine_false(self, platform):
        platform.set_machine(quantum_machine)
        data = platform.qcis_check_regular('H Q0\nH Q2\nCZ Q0 Q2\nM Q0\nM Q2')
        platform.set_machine(sim_machine)
        print(f'data: {data}')
        assert data is False

    def test_qcis_check_regular_quantum_machine_true(self, platform):
        platform.set_machine(quantum_machine)
        data = platform.qcis_check_regular('H Q0\nH Q7\nCZ Q0 Q7\nM Q0\nM Q7')
        platform.set_machine(sim_machine)
        print(f'data: {data}')
        assert data

    def test_get_experiment_circuit(self, platform, lab_id, qcis):
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=qcis,
                                          language=QuantumLanguage.QCIS)
        print(f'exp_id: {exp_id}')
        query_id = platform.run_experiment(exp_id=exp_id, num_shots=3000)
        print(f'query_id: {query_id}')
        data = platform.get_experiment_circuit(query_id)
        print(f'circuit: {data}')
        assert data

    def test_get_experiment_circuit_isq(self, platform, lab_id):
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=self.isq_circuit,
                                          language=QuantumLanguage.ISQ)
        print(f'exp_id: {exp_id}')
        query_id = platform.run_experiment(exp_id=exp_id, num_shots=3000)
        print(f'query_id: {query_id}')
        data = platform.get_experiment_circuit(query_id)
        print(f'circuit: {data}')
        assert data

    def test_re_execute_task(self, platform, lab_id, qcis):
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=qcis,
                                          language=QuantumLanguage.QCIS)
        print(f'exp_id: {exp_id}')
        query_id = platform.run_experiment(exp_id=exp_id, num_shots=3000)
        print(f'query_id: {query_id}')
        data = platform.re_execute_task(query_id=query_id)
        print(f'data: {data}')
        assert data

    def test_re_execute_task_lab_id_query_id(self, platform, lab_id, qcis):
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=qcis,
                                          language=QuantumLanguage.QCIS)
        print(f'exp_id: {exp_id}')
        query_id = platform.run_experiment(exp_id=exp_id, num_shots=3000)
        print(f'query_id: {query_id}')
        data = platform.re_execute_task(lab_id=lab_id, query_id=query_id)
        print(f'data: {data}')
        assert data

    def test_re_execute_task_lab_id(self, platform, lab_id, qcis):
        exp_id_1 = platform.save_experiment(lab_id=lab_id, circuit=qcis,
                                            language=QuantumLanguage.QCIS)
        print(f'exp_id 1: {exp_id_1}')
        query_id_1 = platform.run_experiment(exp_id=exp_id_1, num_shots=3000)
        print(f'query_id 1: {query_id_1}')

        exp_id_2 = platform.save_experiment(lab_id=lab_id, circuit=qcis,
                                            language=QuantumLanguage.QCIS)
        print(f'exp_id 2: {exp_id_2}')
        query_id_2 = platform.run_experiment(exp_id=exp_id_2, num_shots=3000)
        print(f'query_id 2: {query_id_2}')
        data = platform.re_execute_task(lab_id=lab_id)
        print(f'data: {data}')
        assert data

    def test_submit_experiment_lab_id(self, platform, lab_id):
        circuit = "RX Q0 {theta}\nM Q0"
        query_id = platform.submit_experiment(
            circuit=circuit,
            language=QuantumLanguage.QCIS,
            parameters=[['theta']],
            values=[[1.2]],
            lab_id=lab_id,
            num_shots=200,
        )
        assert query_id

    def test_submit_experiment_parameters(self, platform):
        circuit = "RX Q0 {theta}\nM Q0"
        query_id = platform.submit_experiment(
            circuit=circuit,
            language=QuantumLanguage.QCIS,
            parameters=[['theta']],
            values=[[1.2]],
            num_shots=200,
        )
        assert query_id

    def test_submit_experiment_single(self, platform):
        circuit = "RX Q0 1.1\nM Q0"
        query_id = platform.submit_experiment(
            circuit=circuit,
            language=QuantumLanguage.QCIS,
            num_shots=200,
        )
        assert query_id

    def test_submit_experiment_isq_single(self, platform):
        circuit = self.isq_circuit
        query_id = platform.submit_experiment(
            circuit=circuit,
            language=QuantumLanguage.ISQ,
            num_shots=200,
        )
        assert query_id

    def test_submit_experiment_list(self, platform):
        circuit = ["RX Q0 1.1\nM Q0", 'H Q0\nM Q0']
        query_id = platform.submit_experiment(
            circuit=circuit,
            language=QuantumLanguage.QCIS,
            num_shots=200,
        )
        assert query_id

    def test_submit_experiment_name(self, platform):
        circuit = ["RX Q0 1.1\nM Q0", 'H Q0\nM Q0']
        query_id = platform.submit_experiment(
            circuit=circuit,
            language=QuantumLanguage.QCIS,
            num_shots=200,
            name=f'test.{datetime.now().strftime("%Y%m%d%H%M%S")}'
        )
        assert query_id

    def test_submit_experiment_isq(self, platform):
        circuit = [self.isq_circuit, self.isq_circuit]
        query_id = platform.submit_experiment(
            circuit=circuit,
            language=QuantumLanguage.ISQ,
            num_shots=200,
            name=f'test.{datetime.now().strftime("%Y%m%d%H%M%S")}'
        )
        assert query_id

    def test_submit_experiment_isq_lab(self, platform, lab_id):
        circuit = [self.isq_circuit, self.isq_circuit]
        query_id = platform.submit_experiment(
            circuit=circuit,
            lab_id=lab_id,
            language=QuantumLanguage.ISQ,
            num_shots=200,
            name=f'test.{datetime.now().strftime("%Y%m%d%H%M%S")}'
        )
        assert query_id

    def test_submit_experiment_without_verify(self, platform):
        circuit = ["RX Q0 1.1\nM Q0", 'H Q0\nM Q0']
        query_id = platform.submit_experiment(
            circuit=circuit,
            language=QuantumLanguage.QCIS,
            num_shots=200,
            name=f'test.{datetime.now().strftime("%Y%m%d%H%M%S")}',
            is_verify=False
        )
        assert query_id

    def test_submit_experiment_isq_without_verify(self, platform):
        circuit = [self.isq_circuit, self.isq_circuit]
        query_id = platform.submit_experiment(
            circuit=circuit,
            language=QuantumLanguage.ISQ,
            num_shots=200,
            name=f'test.{datetime.now().strftime("%Y%m%d%H%M%S")}',
            is_verify=False
        )
        assert query_id

    def test_auto_read_calibration(self, platform):
        circuits = ['H Q0\nX Q7\nM Q0\nM Q7', 'H Q0\nX Q7\nM Q0\nM Q7']
        platform.set_machine(quantum_machine)
        query_ids = platform.submit_experiment(
            circuit=circuits,
            num_shots=2000,
        )
        data = platform.query_experiment(query_ids, readout_calibration=True)
        assert len(data) == 2
        assert isinstance(data[0]['probability'], dict)

        config = platform.download_config()
        data = platform.query_experiment(query_ids, readout_calibration=True, machine_config=config)
        assert len(data) == 2

        data = platform.query_experiment(query_ids, readout_calibration=False)
        assert len(data) == 2

        query_ids = platform.submit_experiment(
            circuit=circuits,
            num_shots=2000,
            machine_name=sim_machine,
        )
        data = platform.query_experiment(query_ids, readout_calibration=True)
        assert len(data) == 2


@pytest.fixture(params=[
    [{
        "noise_type": "decoherence",
        "params": [0.3, 0.2, 0.3]
    }],
    [{
        "noise_type": "bit-flip",
        "params": [0.3]
    }],
    [{
        "noise_type": "phase-flip",
        "params": [0.3]
    }],
    [{
        "noise_type": "depolarizing",
        "params": [0.1]
    }],
])
def noise(request):
    return request.param


@pytest.fixture(params=[1, 2, 3])
def quantum_state(request):
    return request.param


class TestTianYanSim:
    """
    Test Sim Update
    """

    @staticmethod
    def _run_save_experiment(platform, lab_id, circuit, **kwargs):
        exp_id = platform.save_experiment(lab_id, circuit=circuit, **kwargs)
        query_id = platform.run_experiment(exp_id)
        assert query_id != 0
        data = platform.query_experiment(query_id)
        assert isinstance(data[0]['probability'], dict)

    @staticmethod
    def _run_submit_job(platform, lab_id, circuit, **kwargs):
        query_id = platform.submit_job(lab_id=lab_id, circuit=circuit, **kwargs)
        assert query_id != 0
        data = platform.query_experiment(query_id)
        assert isinstance(data[0]['probability'], dict)

    def test_full_amplitude_sim_save_experiment(self, platform, lab_id, qcis):
        platform.set_machine('zdxlz_simulator')
        self._run_save_experiment(platform, lab_id, qcis)

    def test_stabilizer_save_experiment(self, platform, lab_id, qcis):
        platform.set_machine('stabilizer')
        self._run_save_experiment(platform, lab_id, qcis)

    def test_tn_sim_save_experiment(self, platform, lab_id, qcis):
        platform.set_machine('tensor_network')
        self._run_save_experiment(platform, lab_id, qcis)

    def test_noise_sim_save_experiment(self, platform, lab_id, qcis, noise):
        platform.set_machine('tensor_network_noise')
        self._run_save_experiment(platform, lab_id, qcis, noise=noise)

    def test_single_amplitude_sim_save_experiment(self, platform, lab_id, qcis, quantum_state):
        platform.set_machine('single_amplitude')
        exp_id = platform.save_experiment(lab_id, circuit=qcis, quantum_state=quantum_state)
        query_id = platform.run_experiment(exp_id)
        assert query_id != 0
        data = platform.query_experiment(query_id)
        assert isinstance(data[0]['phase'], str)

    def test_full_amplitude_sim_submit_job(self, platform, lab_id, qcis):
        platform.set_machine('zdxlz_simulator')
        self._run_submit_job(platform=platform, lab_id=lab_id, circuit=qcis)

    def test_stabilizer_submit_job(self, platform, lab_id, qcis):
        platform.set_machine('stabilizer')
        self._run_submit_job(platform, lab_id, qcis)

    def test_tn_sim_submit_job(self, platform, lab_id, qcis):
        platform.set_machine('tensor_network')
        self._run_submit_job(platform, lab_id, qcis)

    def test_noise_sim_save_submit_job(self, platform, lab_id, qcis, noise):
        platform.set_machine('tensor_network_noise')
        self._run_submit_job(platform=platform, lab_id=lab_id, circuit=qcis, noise=noise)

    def test_single_amplitude_sim_submit_job(self, platform, lab_id, qcis, quantum_state):
        platform.set_machine('single_amplitude')
        query_id = platform.submit_job(lab_id=lab_id, circuit=qcis, quantum_state=quantum_state)
        assert query_id != 0
        data = platform.query_experiment(query_id)
        assert isinstance(data[0]['phase'], str)
